/*
 * Copyright (c) 2019 https://github.com/arionik/avilib
 * Copyright (c) 2020-2021 https://gitee.com/fsfzp888
 * All rights reserved
 */

#include <array>
#include <cstdint>
#include <cstring>
#include <deque>
#include <fstream>
#include <iostream>
#include <map>
#include <string>
#include <vector>

struct avilib_BITMAPINFO
{
    uint32_t biSize;
    uint32_t biWidth;
    uint32_t biHeight;
    uint16_t biPlanes;
    uint16_t biBitCount;
    uint32_t biCompression;
    uint32_t biSizeImage;
    uint32_t biXPelsPerMeter;
    uint32_t biYPelsPerMeter;
    uint32_t biClrUsed;
    uint32_t biClrImportant;
};

struct avilib_WAVEFORMATEX
{
    uint16_t wFormatTag;
    uint16_t nChannels;
    uint32_t nSamplesPerSec;
    uint32_t nAvgBytesPerSec;
    uint16_t nBlockAlign;
    uint16_t wBitsPerSample;
};

enum avilib_streamtype_t
{
    avilib_UnknownStreamtype,
    avilib_Video,
    avilib_Audio
};

namespace avilib
{
typedef uint32_t DWORD;
union FOURCC
{
    void operator=(uint32_t i) { _ = i; };
    DWORD _;
    char fcc[4];
};
typedef uint16_t WORD;
typedef int32_t LONG;
typedef uint8_t BYTE;

struct AVIMAINHEADER
{
    FOURCC fcc;
    DWORD cb;
    DWORD dwMicroSecPerFrame;
    DWORD dwMaxBytesPerSec;
    DWORD dwPaddingGranularity;
    DWORD dwFlags;
#define AVIF_HASINDEX 0x00000010  // Index at end of file?
#define AVIF_MUSTUSEINDEX 0x00000020
#define AVIF_ISINTERLEAVED 0x00000100
#define AVIF_TRUSTCKTYPE 0x00000800  // Use CKType to find key frames
#define AVIF_WASCAPTUREFILE 0x00010000
#define AVIF_COPYRIGHTED 0x00020000
    DWORD dwTotalFrames;
    DWORD dwInitialFrames;
    DWORD dwStreams;
    DWORD dwSuggestedBufferSize;
    DWORD dwWidth;
    DWORD dwHeight;
    DWORD dwReserved[4];
};

struct AVISTREAMHEADER
{
    FOURCC fcc;
    DWORD cb;
    FOURCC fccType;
    FOURCC fccHandler;
    DWORD dwFlags;
#define AVISF_DISABLED 0x00000001
#define AVISF_VIDEO_PALCHANGES 0x00010000
    WORD wPriority;
    WORD wLanguage;
    DWORD dwInitialFrames;
    DWORD dwScale;
    DWORD dwRate;
    DWORD dwStart;
    DWORD dwLength;
    DWORD dwSuggestedBufferSize;
    DWORD dwQuality;
    DWORD dwSampleSize;
    struct
    {
        short int left;
        short int top;
        short int right;
        short int bottom;
    } rcFrame;
};

struct BITMAPINFO
{
    // struct tagBITMAPINFO
    //{
    DWORD biSize;
    LONG biWidth;
    LONG biHeight;
    WORD biPlanes;
    WORD biBitCount;
    FOURCC biCompression;
#define BI_RGB 0L
    DWORD biSizeImage;
    LONG biXPelsPerMeter;
    LONG biYPelsPerMeter;
    DWORD biClrUsed;
    DWORD biClrImportant;
    //} bmInfo;
    // struct tagRGBQUAD {
    //	BYTE rgbBlue;
    //	BYTE rgbGreen;
    //	BYTE rgbRed;
    //	BYTE rgbReserved;
    //} bmiColors[1];
};

struct WAVEFORMATEX
{
#define WAVE_FORMAT_EXTENSIBLE 0xfffe
#define WAVE_FORMAT_MPEG 0x0050 /* MPEG Layer 1,2 */
#define WAVE_FORMAT_MPEGLAYER3 0x0055 /* MPEG Layer 3 */
    WORD wFormatTag;
    WORD nChannels;
    DWORD nSamplesPerSec;
    DWORD nAvgBytesPerSec;
    WORD nBlockAlign;
    WORD wBitsPerSample;
    WORD cbSize;
};

struct WAVEFORMATEXTENSIBLE
{
    WAVEFORMATEX Format;
    union
    {
        WORD wValidBitsPerSample;
        WORD wSamplesPerBlock;
        WORD wReserved;
    } Samples;
    DWORD dwChannelMask;
    uint8_t SubFormat[16];
};

struct MPEG1WAVEFORMAT
{
    WAVEFORMATEX wfx;
    WORD fwHeadLayer;
    DWORD dwHeadBitrate;
    WORD fwHeadMode;
    WORD fwHeadModeExt;
    WORD wHeadEmphasis;
    WORD fwHeadFlags;
    DWORD dwPTSLow;
    DWORD dwPTSHigh;
};

struct MPEGLAYER3WAVEFORMAT
{
    WAVEFORMATEX wfx;
    WORD wID;
    DWORD fdwFlags;
    WORD nBlockSize;
    WORD nFramesPerBlock;
    WORD nCodecDelay;
};

struct ODMLExtendedAVIHeader
{
    DWORD dwTotalFrames;
};

struct AVIOLDINDEX
{
    FOURCC dwChunkId;
    DWORD dwFlags;
#define AVIIF_LIST 0x00000001
#define AVIIF_KEYFRAME 0x00000010
#define AVIIF_NO_TIME 0x00000100
    DWORD dwOffset;
    DWORD dwSize;
};

#define AVI_INDEX_OF_INDEXES 0x00
#define AVI_INDEX_OF_CHUNKS 0x01
#define AVI_INDEX_OF_TIMED_CHUNKS 0x02
#define AVI_INDEX_OF_SUB_2FIELD 0x03
#define AVI_INDEX_IS_DATA 0x80

// index subtype codes
//
#define AVI_INDEX_SUB_DEFAULT 0x00

// INDEX_OF_CHUNKS subtype codes
//
#define AVI_INDEX_SUB_2FIELD 0x01

struct AVISTDINDEX
{
    FOURCC fcc;  // "ix##"
    DWORD cb;
    WORD wLongsPerEntry;    // must be sizeof(aIndex[0])/sizeof(DWORD)
    BYTE bIndexSubType;     // must be 0
    BYTE bIndexType;        // must be AVI_INDEX_OF_CHUNKS
    DWORD nEntriesInUse;    //
    FOURCC dwChunkId;       // "##dc" or "##db" or "##wb", etc.
    uint64_t qwBaseOffset;  // all dwOffsets in aIndex array are relative to this
    DWORD dwReserved;       // must be 0
};

struct AVISTDINDEX_ENTRY
{
    DWORD dwOffset;  // qwBaseOffset + this is absolute file offset
    DWORD dwSize;    // bit 31 is set if this is NOT a keyframe
};

struct AVISUPERINDEX
{
    FOURCC fcc;           // 'xdni'
    DWORD cb;             // size of this structure
    WORD wLongsPerEntry;  // must be 4 (size of each entry in aIndex array)
    BYTE bIndexSubType;   // must be 0 or AVI_INDEX_2FIELD
    BYTE bIndexType;      // must be AVI_INDEX_OF_INDEXES
    DWORD nEntriesInUse;  // number of entries in aIndex array that are used
    FOURCC dwChunkId;     // "##dc" or "##db" or "##wb", etc.
    DWORD dwReserved[3];  // must be 0
};

struct AVISUPERINDEX_ENTRY
{
    uint64_t qwOffset;  // absolute file offset, offset 0 is unused entry??
    DWORD dwSize;       // size of index chunk at this offset
    DWORD dwDuration;   // time span in stream ticks
};

struct VIDEO_PROP_HEADER
{
    enum video_format_t
    {
        FORMAT_UNKNOWN,
        FORMAT_PAL_SQUARE,
        FORMAT_PAL_CCIR_601,
        FORMAT_NTSC_SQUARE,
        FORMAT_NTSC_CCIR_601
    };
    DWORD VideoFormatToken;
    enum video_standard_t
    {
        STANDARD_UNKNOWN,
        STANDARD_PAL,
        STANDARD_NTSC,
        STANDARD_SECAM
    };
    DWORD VideoStandard;
    DWORD dwVerticalRefreshRate;
    DWORD dwHTotalInT;
    DWORD dwVTotalInLines;
    DWORD dwFrameAspectRatio;
    DWORD dwFrameWidthInPixels;
    DWORD dwFrameHeightInLines;
    DWORD nbFieldPerFrame;
    struct VIDEO_FIELD_DESC
    {
        DWORD CompressedBMHeight;
        DWORD CompressedBMWidth;
        DWORD ValidBMHeight;
        DWORD ValidBMWidth;
        DWORD ValidBMXOffset;
        DWORD ValidBMYOffset;
        DWORD VideoXOffsetInT;
        DWORD VideoYValidStartLine;
    };
    // VIDEO_FIELD_DESC FieldInfo[nbFieldPerFrame];
};

int32_t gcd(int32_t nom, int32_t denom);
bool cancel(int32_t nom, int32_t denom, int32_t &out_nom, int32_t &out_denom);

class LoggingObject
{
    std::deque<std::string> m_messages;

  protected:
    void log(const std::string &);
    void clear_log();

  public:
    bool getNextMessage(std::string &message);
};

FOURCC generate_fcc(std::string id, uint8_t idx);

class AviWriter : protected LoggingObject
{
    typedef int32_t stream_id;

    std::ofstream _f;
    uint64_t pos_RIFFSize;
    uint64_t pos_moviListSize;
    std::ofstream::streamoff pos_odmlExt;
    std::ofstream::streamoff pos_1stMoviStart;
    std::ofstream::streamoff pos_aviMainHeader;
    std::map<stream_id, std::ofstream::streamoff> pos_odmlSuperIdx;
    std::map<stream_id, std::ofstream::streamoff> pos_streamHeader;

    std::map<stream_id, avilib::AVISUPERINDEX> m_superIdxs;
    std::map<stream_id, std::vector<avilib::AVISTDINDEX>> m_stdIndexes;
    std::map<stream_id, std::map<int32_t, std::deque<avilib::AVISTDINDEX_ENTRY>>> m_stdIndexEntries;
    std::map<stream_id, avilib_streamtype_t> m_streamTypes;

    std::deque<avilib::AVIOLDINDEX> m_oldIndexEntries;

    avilib::AVIMAINHEADER m_avimHeader;
    std::map<int, avilib::AVISTREAMHEADER> m_avisHeaders;
    avilib::BITMAPINFO m_bitmapInfo;
    avilib::WAVEFORMATEXTENSIBLE m_waveFormat;

    bool m_openDML             = true;
    bool m_opened              = false;
    uint32_t m_videoStreamIdx  = ~0u;
    uint32_t m_reservedFrames  = 0;
    uint32_t m_RIFF_idx        = 0;  // nth RIFF section (OpenDML)
    uint32_t m_RIFF_size       = 0;  // current RIFF size (OpenDML)
    uint32_t m_totalFrames     = 0;
    uint64_t m_currBaseOff     = 0;
    uint64_t m_currMaxRIFFSize = 1000 * 1000 * 1000;

    void finishOld();
    void finishAvix();
    void reset();

  public:
    AviWriter();
    virtual ~AviWriter();

    bool setVideoProperties(uint8_t stream_idx, int32_t width, int32_t height, uint32_t codec, uint32_t framesize, double rateHz);
    bool setAudioProperties(uint8_t stream_idx,
                            int16_t format,
                            int8_t channels,
                            uint32_t samplesPerSecond,
                            uint32_t avgBytesPerSecond,
                            uint16_t bitsPerSample,
                            uint16_t blockAlign);
    bool setFrameCount(uint8_t stream_idx, uint32_t count);

    bool open(const char *filename);
    bool writeFrame(uint8_t stream_idx, void *data);
    bool writeData(uint8_t stream_idx, void *data, uint32_t len);
    bool close();
    bool openDML(bool yes);
};

class AviReader : protected LoggingObject
{
    std::ifstream _f;

    typedef struct _dmlindex
    {
        FOURCC dwChunkId;
        uint32_t u32_flags;
        int64_t i64_offset;
        uint32_t u32_size;
    } DMLINDEX;

    // AVI header stuff
    avilib::AVIMAINHEADER m_avimHeader;
    avilib::AVISTREAMHEADER m_avisHeader[2];
    std::map<int32_t, avilib::AVISUPERINDEX> m_superIndex;
    avilib::BITMAPINFO m_bitmapInfo;
    avilib::WAVEFORMATEX m_waveformat;
    avilib::WAVEFORMATEXTENSIBLE m_waveformatExt;
    avilib::MPEG1WAVEFORMAT m_waveformatMpeg12;
    avilib::MPEGLAYER3WAVEFORMAT m_waveformatMP3;

    avilib::ODMLExtendedAVIHeader m_odmlExt;

    std::map<int32_t, std::deque<DMLINDEX>> m_frameIdxs;
    std::map<int32_t, std::deque<DMLINDEX>> odml_frameIdxs;

    std::deque<uint64_t> m_moviOffs;
    int32_t m_videoStreamIdx = 0;
    bool m_useMovieOffset    = false;

  public:
    AviReader();
    virtual ~AviReader();

    bool open(const char *filename);
    bool close();
    int32_t readFrame(uint32_t idx, uint32_t stream, void *data);
    bool getSize(int32_t &width, int32_t &height);
    bool getFrameRate(double &rateHz);
    bool getFormat(uint32_t stream, avilib_BITMAPINFO &bm_info);
    bool getFormat(uint32_t stream, avilib_WAVEFORMATEX &wav_info);
    bool getFrameCount(uint32_t stream, uint32_t &count);
    bool getCodec(uint32_t stream, uint32_t &codec);
    bool getAllocSize(uint32_t stream, uint32_t &size);
    bool getStreamCount(uint32_t &count);
    bool getStreamType(uint32_t stream, avilib_streamtype_t &type);
};

}  // namespace avilib
